package adkr

import (
	"ebbflow/pkg/crypto"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"math/big"
)

// ProofOfKnowledgeData represents a zero-knowledge proof of discrete logarithm
type ProofOfKnowledgeData struct {
	T  crypto.G1           `json:"t"`  // T = g^u
	Z  crypto.ZR           `json:"z"`  // z = u + e*x
	EU *PaillierCiphertext `json:"eu"` // encrypted u
	W  *big.Int            `json:"w"`  // w = r^e * s mod n^2
}

// PaillierCiphertext represents a Paillier encrypted value
type PaillierCiphertext struct {
	C *big.Int `json:"c"`
}

// ProveKnowledgeOfEncryptedDLog generates a proof of knowledge of discrete logarithm
// This corresponds to the Python function prove_knowledge_of_encrypted_dlog_BN
func ProveKnowledgeOfEncryptedDLog(
	g crypto.G1,
	x crypto.ZR,
	pk *crypto.PaillierPublicKey,
	gToTheX *crypto.G1,
) ([]interface{}, error) {
    // Fast-bench mode or missing pk: return a lightweight, deterministic dummy proof
    if IsFastBench() || pk == nil {
        // layout compatible with Verify: [TBytes, z, eu, w]
        // We just return zeros; verifier will short-circuit in fast mode.
        var zeroZ crypto.ZR
        zeroZ.SetInt64(0)
        dummy := []interface{}{serializeG1(&g), zeroZ, &PaillierCiphertext{C: big.NewInt(0)}, big.NewInt(0)}
        return dummy, nil
    }
	var Y crypto.G1
	if gToTheX == nil {
		Y = crypto.MulG1(&g, &x)
	} else {
		Y = *gToTheX
	}

	// Generate random values
	r := pk.GetRandomLtN()
	u := new(big.Int).Div(pk.GetRandomLtN(), big.NewInt(3)) // maximum valid value we can encrypt

	// Encrypt x
	c, err := pk.Encrypt(ZRToBigInt(&x), r)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt x: %w", err)
	}

	// Compute T = g^u
	var uZr crypto.ZR
	uZr.SetBigInt(u)
	T := crypto.MulG1(&g, &uZr)

	// Compute challenge e = H(pk, g, Y, c, T)
	e := computeChallenge(pk, &g, &Y, c.C, &T)

	// Compute z = u + e*x
	var eZr crypto.ZR
	eZr.SetBigInt(e)
	var z crypto.ZR
	z.Mul(&eZr, &x)
	z.Add(&z, &uZr)

	// Encrypt u
	s := pk.GetRandomLtN()
	eu, err := pk.Encrypt(u, s)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt u: %w", err)
	}

	// Compute w = r^e * s mod n^2
	w := new(big.Int).Exp(r, e, pk.NSquare)
	w.Mul(w, s)
	w.Mod(w, pk.NSquare)

	_ = ProofOfKnowledgeData{
		T:  T,
		Z:  z,
		EU: &PaillierCiphertext{C: eu.C},
		W:  w,
	}

	// Return in the format expected by the Python code
	return []interface{}{
		serializeG1(&Y),
		c.C,
		[]interface{}{serializeG1(&T), z, eu, w},
	}, nil
}

// VerifyKnowledgeOfDiscreteLog verifies a proof of knowledge of discrete logarithm
// This corresponds to the Python function verify_knowledge_of_discrete_log_BN
func VerifyKnowledgeOfDiscreteLog(
	pk *crypto.PaillierPublicKey,
	g crypto.G1,
	Y []byte,
	c *big.Int,
	proof []interface{},
) bool {
    // Fast-bench mode: skip verification
    if IsFastBench() {
        return true
    }
	if len(proof) != 4 {
		return false
	}

	TBytes, ok1 := proof[0].([]byte)
	z, ok2 := proof[1].(crypto.ZR)
	eu, ok3 := proof[2].(*PaillierCiphertext)
	w, ok4 := proof[3].(*big.Int)

	if !ok1 || !ok2 || !ok3 || !ok4 {
		return false
	}

	// Deserialize T
	T, err := deserializeG1(TBytes)
	if err != nil {
		return false
	}

	// Compute challenge e = H(pk, g, Y, c, T)
	e := computeChallenge(pk, &g, Y, c, &T)

	// Check T == g^z * Y^(-e)
	var eZr crypto.ZR
	eZr.SetBigInt(e)
	var negE crypto.ZR
	negE.Neg(&eZr)

	expectedT1 := crypto.MulG1(&g, &z)
	expectedT2 := crypto.MulG1(deserializeG1FromBytes(Y), &negE)
	expectedT := crypto.AddG1(&expectedT1, &expectedT2)

	if !T.G1Affine.Equal(&expectedT.G1Affine) {
		return false
	}

	// Check (eu.ciphertext * c^e) mod n^2 == pk.encrypt(z, w).ciphertext
	cE := new(big.Int).Exp(c, e, pk.NSquare)
	leftSide := new(big.Int).Mul(eu.C, cE)
	leftSide.Mod(leftSide, pk.NSquare)

	rightSide, err := pk.Encrypt(ZRToBigInt(&z), w)
	if err != nil {
		return false
	}

	return leftSide.Cmp(rightSide.C) == 0
}

// computeChallenge computes the challenge e = H(pk, g, Y, c, T)
func computeChallenge(pk *crypto.PaillierPublicKey, g *crypto.G1, Y interface{}, c *big.Int, T *crypto.G1) *big.Int {
	// Serialize all components
	pkBytes, _ := json.Marshal(pk)
	gBytes := serializeG1(g)
	var yBytes []byte
	switch y := Y.(type) {
	case []byte:
		yBytes = y
	case *crypto.G1:
		yBytes = serializeG1(y)
	default:
		yBytes = []byte(fmt.Sprintf("%v", y))
	}
	cBytes := c.Bytes()
	tBytes := serializeG1(T)

	// Compute hash
	h := sha256.New()
	h.Write(pkBytes)
	h.Write(gBytes)
	h.Write(yBytes)
	h.Write(cBytes)
	h.Write(tBytes)
	hash := h.Sum(nil)

	// Convert to big.Int
	return new(big.Int).SetBytes(hash)
}

// serializeG1 serializes a G1 point to bytes
func serializeG1(g *crypto.G1) []byte {
	// Use the existing MarshalJSON method
	bytes, _ := g.MarshalJSON()
	return bytes
}

// deserializeG1 deserializes bytes to a G1 point
func deserializeG1(data []byte) (crypto.G1, error) {
	var g crypto.G1
	err := g.UnmarshalJSON(data)
	return g, err
}

// deserializeG1FromBytes deserializes bytes to a G1 point (helper function)
func deserializeG1FromBytes(data []byte) *crypto.G1 {
	g, err := deserializeG1(data)
	if err != nil {
		// Return identity element on error
		return &crypto.G1{}
	}
	return &g
}

// BigInt returns the big.Int representation of a ZR element
func ZRToBigInt(z *crypto.ZR) *big.Int {
	return z.Element.BigInt(new(big.Int))
}
